"""
Copyright 2020 Huawei Technologies Co., Ltd

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""

import os
import argparse
import traceback
from common import set_env
from common import atc_model
from flat_generate_model import get_dist_compute_json
from sq8_generate_model import get_distance_compute_sq8_json

# number of code per time
code_num = 65536


def arg_parse():
    """
    Parse arguements to the operator model
    """

    parser = argparse.ArgumentParser(
        description='generate distance_compute_sq8 operator model')

    parser.add_argument("-d",
                        dest='dim',
                        default=128,
                        type=int,
                        help="Feature dimension")
    parser.add_argument("-c",
                        dest='coarse_centroid_num',
                        default=16384,
                        type=int,
                        help="Number of coarse centroid")

    return parser.parse_args()


def generate_by_atc():
    # generate directory of model
    dirs = './op_models/'
    if not os.path.exists(dirs):
        os.makedirs(dirs)

    atc_model('distance_compute_op1')
    atc_model('distance_compute_op2')
    atc_model('distance_compute_op4')
    atc_model('distance_compute_op8')
    atc_model('distance_compute_op16')
    atc_model('distance_compute_op32')
    atc_model('distance_compute_op64')
    atc_model('distance_compute_sq8_op')


if __name__ == '__main__':
    args = arg_parse()
    set_env()

    dim = args.dim
    coarse_centroid_num = args.coarse_centroid_num

    # generate directory of config
    dirs = './config/'
    if not os.path.exists(dirs):
        os.makedirs(dirs)

    get_dist_compute_json(1, dim, coarse_centroid_num,
                          './config/distance_compute_op1.json')
    get_dist_compute_json(2, dim, coarse_centroid_num,
                          './config/distance_compute_op2.json')
    get_dist_compute_json(4, dim, coarse_centroid_num,
                          './config/distance_compute_op4.json')
    get_dist_compute_json(8, dim, coarse_centroid_num,
                          './config/distance_compute_op8.json')
    get_dist_compute_json(16, dim, coarse_centroid_num,
                          './config/distance_compute_op16.json')
    get_dist_compute_json(32, dim, coarse_centroid_num,
                          './config/distance_compute_op32.json')
    get_dist_compute_json(64, dim, coarse_centroid_num,
                          './config/distance_compute_op64.json')

    get_distance_compute_sq8_json(
        1, dim, code_num, './config/distance_compute_sq8_op.json')

    try:
        generate_by_atc()
    except OSError:
        traceback.print_exc()
        print("make sure set right ASCEND_HOME environment variable")
    else:
        traceback.print_exc()
